import secrets
import time
from typing import Any

from mcp.shared.auth import OAuthClientInformationFull, OAuthToken
from mcp.server.auth.provider import (
    OAuthAuthorizationServerProvider,
    AuthorizationParams,
    AuthorizationCode,
    RefreshToken,
    AccessToken
)
class GZYOAuthServerProvider(OAuthAuthorizationServerProvider):
    def __init__(self, login_url: str):
        self.login_url = login_url # 要传分路由的
        self.clients: dict[str, OAuthClientInformationFull] = {}
        self.state_mapping: dict[str, Any] = {}
        self.auth_codes: dict[str, AuthorizationCode] = {} # 授权码
        self.tokens: dict[str, AccessToken] = {}

    async def get_client(self, client_id: str) -> OAuthClientInformationFull | None:
        """
        客户端 ID 检索客户端信息
        :param client_id: 客户端 ID
        :return: 客户端信息
        """
        return self.clients.get(client_id)

    async def register_client(self, client_info: OAuthClientInformationFull) -> None:
        """
        保存客户端注册的信息
        :param client_info: 客户端信息
        :return: 客户端信息存储到self.clients的字典里
        """
        self.clients[client_info.client_id] = client_info

    async def authorize(self, client: OAuthClientInformationFull, params: AuthorizationParams) -> str:
        """
        作为 /authorize 端点的一部分被调用，并返回一个客户端将被重定向到的 URL。
        许多 MCP 实现会重定向到第三方提供商，以便与该提供商进行第二次 OAuth 交换。在这种设置中，客户端与 MCP 服务器建立了 OAuth 连接，
        而 MCP 服务器与第三方提供商建立了 OAuth 连接。在此流程结束时，客户端应被重定向到 params.redirect_uri 中的 redirect_uri。
         +--------+     +------------+     +-------------------+
        |        |     |            |     |                   |
        | Client | --> | MCP Server | --> | 3rd Party OAuth   |
        |        |     |            |     | Server            |
        +--------+     +------------+     +-------------------+
                            |   ^                  |
        +------------+      |   |                  |
        |            |      |   |    Redirect      |
        |redirect_uri|<-----+   +------------------+
        |            |
        +------------+
        实现需要在 MCP 服务器返回流程上定义另一个处理程序，以执行第二次重定向，并在完成 OAuth 授权步骤时生成和存储授权码。
        :param client: 请求授权的客户端
        :param params: 授权请求的参数
        :return: 用于将客户端重定向到授权的 URL
        """
        # /authorize并不会验证用户的账号和密码，而是记录用户的一些信息，然后返回一个重定向的URL
        # 浏览器会自动跳转到重定向的URL上去
        # authorize只是简单做一些信息的记录
        # state：是由客户端生成的随机的字符串
        state = params.state
        # 以下这种数据一般是存储在redis或者memCache缓存数据库中，Memcached 是运行在NodeJs服务器端的库
        self.state_mapping[state] = {
            "client_id": client.client_id,
            "redirect_uri": params.redirect_uri,
            "state": params.state,
            "code_challenge": params.code_challenge,
            "resource": params.resource,
            "redirect_uri_provided_explicitly": params.redirect_uri_provided_explicitly,
        }
        # 返回重定向的URL
        return f"{self.login_url}?state={state}&client_id={client.client_id}"

    async def load_authorization_code(
            self, client: OAuthClientInformationFull, authorization_code: str
    ) -> AuthorizationCode | None:
        """
        通过授权码加载 AuthorizationCode对象
        :param client: 请求授权码的客户端
        :param authorization_code: 授权码
        :return: 对应的 AuthorizationCode，若未找到则返回 None
        """
        return self.auth_codes.get(authorization_code)

    async def exchange_authorization_code(
            self, client: OAuthClientInformationFull, authorization_code: AuthorizationCode
    ) -> OAuthToken:
        """
        通过authorization_code换取访问令牌和刷新令牌
        :param client: 交换授权码的客户端
        :param authorization_code: 授权码
        :return: OAuth 令牌，包含访问令牌和刷新令牌
        """
        if authorization_code.code not in self.auth_codes:
            raise ValueError("授权码不可用！") # 抛异常
        # 生成access_token，JWT
        token = f"mcp_{secrets.token_hex(32)}"
        # 将生成的token存储在服务器，下次进行校验的时候，就可以从中获取，来判断是否存在，以及是否过期
        self.tokens[token] = AccessToken(
            token=token,
            client_id=client.client_id,
            scopes=['user'],
            expires_at=int(time.time()) + 3600, # 60分钟
            resource=authorization_code.resource
        )
        # 生成完了access_token后，授权码就可以删掉了
        del self.auth_codes[authorization_code.code]
        return OAuthToken(
            access_token=token,
            expires_in=3600,
            scope="user"
        )

    async def load_refresh_token(self, client: OAuthClientInformationFull, refresh_token: str) -> RefreshToken | None:
        """
        通过刷新令牌字符串加载刷新令牌
        :param client: 客户端
        :param refresh_token: 要加载的刷新令牌字符串
        :return: 若找到则返回 RefreshToken 对象，若未找到则返回 None
        """
        raise NotImplementedError("refresh token没有实现！")

    async def exchange_refresh_token(
            self,
            client: OAuthClientInformationFull,
            refresh_token: RefreshToken,
            scopes: list[str],
    ) -> OAuthToken:
        raise NotImplementedError("refresh token没有实现！")

    async def load_access_token(self, token: str) -> AccessToken | None:
        """
        通过令牌加载访问令牌
        :param token: 要验证的访问令牌
        :return: AuthInfo，若令牌无效则返回 None
        """
        access_token = self.tokens.get(token)
        # 判断token是否过期
        if access_token and access_token.expires_at < time.time():
            del self.tokens[token]
            return None

        return access_token

    async def revoke_token(
            self,
            token: AccessToken | RefreshToken,
    ) -> None:
        """
        撤销访问令牌或刷新令牌,如果给定的令牌无效或已被撤销，此方法应不执行任何操作。
        无论提供的是访问令牌还是刷新令牌，实现都应同时撤销访问令牌及其对应的刷新令牌
        :param token: 要撤销的令牌
        :return:
        """
        if token in self.tokens:
            del self.tokens[token.token]

    async def create_authorization_code(self, state: str) -> AuthorizationCode:
        """
        生成授权码
        :param state: 客户端状态
        :return: AuthorizationCode对象的授权码对象
        """
        # 获取之前执行/authorize时，存储的用户授权数据
        state_data = self.state_mapping.get(state)

        redirect_uri = state_data["redirect_uri"]
        code_challenge = state_data["code_challenge"]
        redirect_uri_provided_explicitly = state_data["redirect_uri_provided_explicitly"]
        client_id = state_data["client_id"]
        resource = state_data.get("resource")  # RFC 8707

        # 生成授权码
        code = f"mcp_{secrets.token_hex(16)}"
        auth_code = AuthorizationCode(
            code=code,
            scopes=['user'],
            expires_at=time.time() + 300,
            client_id=client_id,
            code_challenge=code_challenge,
            redirect_uri=redirect_uri,
            redirect_uri_provided_explicitly=redirect_uri_provided_explicitly,
            resource=resource,
        )
        self.auth_codes[code] = auth_code # 保存
        return auth_code